#include "par_Vector.hpp"

template<typename idx_t, typename data_t, int dof>
par_Vector<idx_t, data_t, dof>::par_Vector(MPI_Comm _comm): comm(_comm)
{
}

template<typename idx_t, typename data_t, int dof>
par_Vector<idx_t, data_t, dof>::par_Vector(MPI_Comm _comm, const idx_t _glb_nrows, const idx_t _my_beg, const idx_t _my_end)
    : comm(_comm), glb_nrows(_glb_nrows), beg_row(_my_beg), end_row(_my_end) 
{
    local_vector = new seq_Vector<idx_t, data_t, dof>(end_row - beg_row);
}

template<typename idx_t, typename data_t, int dof>
par_Vector<idx_t, data_t, dof>::par_Vector(const par_Vector & model)
    : comm(model.comm), glb_nrows(model.glb_nrows), beg_row(model.beg_row), end_row(model.end_row)
{
    local_vector = new seq_Vector<idx_t, data_t, dof>(*(model.local_vector));
}

template<typename idx_t, typename data_t, int dof>
par_Vector<idx_t, data_t, dof>::~par_Vector()
{
    if (local_vector) delete local_vector;
}

template<typename idx_t, typename data_t, int dof>
void par_Vector<idx_t, data_t, dof>::set_val(data_t val)
{
    local_vector->operator=(val);
}

template<typename idx_t, typename data_t, int dof>
void par_Vector<idx_t, data_t, dof>::set_values_hypre(const HYPRE_ParVector hypre_vec)
{
    assert(hypre_vec->global_size == glb_nrows);
    assert(hypre_vec->first_index == beg_row);
    assert(hypre_vec->last_index  == end_row - 1);

    #pragma omp parallel for schedule(static)
    for (long long i = 0; i < local_vector->tot_len; i++) {
        local_vector->data[i] = hypre_vec->local_vector->data[i];
    }
}

#ifdef USE_FP16
template class par_Vector<int, __fp16, 1>;
template class par_Vector<int, __fp16, 2>;
template class par_Vector<int, __fp16, 3>;
template class par_Vector<int, __fp16, 4>;
#endif
template class par_Vector<int, float , 1>;
template class par_Vector<int, float , 2>;
template class par_Vector<int, float , 3>;
template class par_Vector<int, float , 4>;

template class par_Vector<int, double, 1>;
template class par_Vector<int, double, 2>;
template class par_Vector<int, double, 3>;
template class par_Vector<int, double, 4>;


template<typename idx_t, typename data_t, typename res_t, int dof>
res_t vec_dot(const par_Vector<idx_t, data_t, dof> & x, const par_Vector<idx_t, data_t, dof> & y)
{
    res_t loc_prod = 0.0, glb_prod = 0.0;
    assert(x.beg_row == y.beg_row && x.end_row == y.end_row);
    
    loc_prod = seq_vec_dot<idx_t, data_t, res_t, dof>(*(x.local_vector), *(y.local_vector));

    // 这里不能用x.comm_pkg->mpi_scalar_type，因为点积希望向上保留精度
    if constexpr (sizeof(res_t) == 8)
        MPI_Allreduce(&loc_prod, &glb_prod, 1, MPI_DOUBLE, MPI_SUM, x.comm);
    else 
        MPI_Allreduce(&loc_prod, &glb_prod, 1, MPI_FLOAT , MPI_SUM, x.comm);
    
    return glb_prod;
}
template double vec_dot(par_Vector<int, double, 1> const &, par_Vector<int, double, 1> const &);
template double vec_dot(par_Vector<int, double, 2> const &, par_Vector<int, double, 2> const &);
template double vec_dot(par_Vector<int, double, 3> const &, par_Vector<int, double, 3> const &);
template double vec_dot(par_Vector<int, double, 4> const &, par_Vector<int, double, 4> const &);
template double vec_dot(par_Vector<int, float , 1> const &, par_Vector<int, float , 1> const &);
template double vec_dot(par_Vector<int, float , 2> const &, par_Vector<int, float , 2> const &);
template double vec_dot(par_Vector<int, float , 3> const &, par_Vector<int, float , 3> const &);
template double vec_dot(par_Vector<int, float , 4> const &, par_Vector<int, float , 4> const &);

template<typename idx_t, typename data_t, typename res_t, int dof>
res_t vec_inf_norm(par_Vector<idx_t, data_t, dof> const & x)
{   
    assert(sizeof(res_t) == 8);

    res_t loc_prod = 0.0, glb_prod;
    long long max_idx = -1;
    loc_prod = seq_vec_inf_norm<idx_t, data_t, res_t>(*(x.local_vector), & max_idx);

    // 这里不能用x.comm_pkg->mpi_scalar_type，因为点积希望向上保留精度
    if constexpr (sizeof(res_t) == 8)
        MPI_Allreduce(&loc_prod, &glb_prod, 1, MPI_DOUBLE, MPI_MAX, x.comm);
    else 
        MPI_Allreduce(&loc_prod, &glb_prod, 1, MPI_FLOAT , MPI_MAX, x.comm);
    
    return glb_prod;
}
template double vec_inf_norm(par_Vector<int, double, 1> const &);
template double vec_inf_norm(par_Vector<int, double, 2> const &);
template double vec_inf_norm(par_Vector<int, double, 3> const &);
template double vec_inf_norm(par_Vector<int, double, 4> const &);
template double vec_inf_norm(par_Vector<int, float , 1> const &);
template double vec_inf_norm(par_Vector<int, float , 2> const &);
template double vec_inf_norm(par_Vector<int, float , 3> const &);
template double vec_inf_norm(par_Vector<int, float , 4> const &);

// v1 + alpha * v2 -> v
template<typename idx_t, typename data_t, typename scalar_t, int dof>
void vec_add(const par_Vector<idx_t, data_t, dof> & v1, scalar_t alpha, 
             const par_Vector<idx_t, data_t, dof> & v2, par_Vector<idx_t, data_t, dof> & v)
{
    assert(v1.beg_row == v2.beg_row && v1.end_row == v2.end_row);
    assert(v1.beg_row == v .beg_row && v1.end_row == v .end_row);

    seq_vec_add(*(v1.local_vector), alpha, *(v2.local_vector), *(v.local_vector));
}
template void vec_add(const par_Vector<int, float , 1> &, float , const par_Vector<int, float , 1> &, par_Vector<int, float , 1> &);
template void vec_add(const par_Vector<int, float , 2> &, float , const par_Vector<int, float , 2> &, par_Vector<int, float , 2> &);
template void vec_add(const par_Vector<int, float , 3> &, float , const par_Vector<int, float , 3> &, par_Vector<int, float , 3> &);
template void vec_add(const par_Vector<int, float , 4> &, float , const par_Vector<int, float , 4> &, par_Vector<int, float , 4> &);
template void vec_add(const par_Vector<int, double, 1> &, double, const par_Vector<int, double, 1> &, par_Vector<int, double, 1> &);
template void vec_add(const par_Vector<int, double, 2> &, double, const par_Vector<int, double, 2> &, par_Vector<int, double, 2> &);
template void vec_add(const par_Vector<int, double, 3> &, double, const par_Vector<int, double, 3> &, par_Vector<int, double, 3> &);
template void vec_add(const par_Vector<int, double, 4> &, double, const par_Vector<int, double, 4> &, par_Vector<int, double, 4> &);

// src -> dst
template<typename idx_t, typename data_t, int dof>
void vec_copy(const par_Vector<idx_t, data_t, dof> & src, par_Vector<idx_t, data_t, dof> & dst)
{
    assert(src.beg_row == dst.beg_row && src.end_row == dst.end_row);

    seq_vec_copy(*(src.local_vector), *(dst.local_vector));
}
template void vec_copy(par_Vector<int, float , 1> const &, par_Vector<int, float , 1> &);
template void vec_copy(par_Vector<int, float , 2> const &, par_Vector<int, float , 2> &);
template void vec_copy(par_Vector<int, float , 3> const &, par_Vector<int, float , 3> &);
template void vec_copy(par_Vector<int, float , 4> const &, par_Vector<int, float , 4> &);
template void vec_copy(par_Vector<int, double, 1> const &, par_Vector<int, double, 1> &);
template void vec_copy(par_Vector<int, double, 2> const &, par_Vector<int, double, 2> &);
template void vec_copy(par_Vector<int, double, 3> const &, par_Vector<int, double, 3> &);
template void vec_copy(par_Vector<int, double, 4> const &, par_Vector<int, double, 4> &);

// coeff * src -> dst
template<typename idx_t, typename data_t, typename scalar_t, int dof>
void vec_mul_by_scalar(const scalar_t coeff, const par_Vector<idx_t, data_t, dof> & src, par_Vector<idx_t, data_t, dof> & dst)
{
    assert(src.beg_row == dst.beg_row && src.end_row == dst.end_row);

    seq_vec_mul_by_scalar(coeff, *(src.local_vector), *(dst.local_vector));
}
template void vec_mul_by_scalar(const float , const par_Vector<int, float , 1> &, par_Vector<int, float , 1> &);
template void vec_mul_by_scalar(const float , const par_Vector<int, float , 2> &, par_Vector<int, float , 2> &);
template void vec_mul_by_scalar(const float , const par_Vector<int, float , 3> &, par_Vector<int, float , 3> &);
template void vec_mul_by_scalar(const float , const par_Vector<int, float , 4> &, par_Vector<int, float , 4> &);
template void vec_mul_by_scalar(const double, const par_Vector<int, double, 1> &, par_Vector<int, double, 1> &);
template void vec_mul_by_scalar(const double, const par_Vector<int, double, 2> &, par_Vector<int, double, 2> &);
template void vec_mul_by_scalar(const double, const par_Vector<int, double, 3> &, par_Vector<int, double, 3> &);
template void vec_mul_by_scalar(const double, const par_Vector<int, double, 4> &, par_Vector<int, double, 4> &);

// vec *= coeff
template<typename idx_t, typename data_t, typename scalar_t, int dof>
void vec_scale(const scalar_t coeff, par_Vector<idx_t, data_t, dof> & vec)
{
    seq_vec_scale(coeff, *(vec.local_vector));
}
template void vec_scale(const float , par_Vector<int, float , 1> &);
template void vec_scale(const float , par_Vector<int, float , 2> &);
template void vec_scale(const float , par_Vector<int, float , 3> &);
template void vec_scale(const float , par_Vector<int, float , 4> &);
template void vec_scale(const double, par_Vector<int, double, 1> &);
template void vec_scale(const double, par_Vector<int, double, 2> &);
template void vec_scale(const double, par_Vector<int, double, 3> &);
template void vec_scale(const double, par_Vector<int, double, 4> &);

